import chromadb
from chromadb.utils import embedding_functions

custom_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

client = chromadb.PersistentClient(path="./chromadb_data")

collection = client.get_or_create_collection(
    "collection3", embedding_function=custom_ef
)

""" collection.add(
    documents=["这是第一个文档", "这是第二个文档", "这是第三个文档"],
    metadatas=[{"source": "doc1"}, {"source": "doc2"}, {"source": "doc3"}],
    ids=["id1", "id2", "id3"],
) """

results = collection.query(
    query_texts=["文档"], n_results=1, where={"source": {"$eq": "doc2"}}
)
print(results)
